# -*- encoding=utf-8 -*-
from __future__ import print_function
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedShuffleSplit
import argparse

parser = argparse.ArgumentParser(description="just split it")
parser.add_argument('--type', type=str,
                    choices=['test', 'train'])
parser.add_argument('--test-list', type=str)
parser.add_argument('--weights', type=str)
parser.add_argument('--arch', type=str)

args = parser.parse_args()

Test = '/home/ping501b/member/gpio/pytorch-coviar/data/YLIMED/YLIMED_short_mp4_EvAll_split_test.txt'
Train = '/home/ping501b/member/gpio/pytorch-coviar/data/YLIMED/YLIMED_short_mp4_EvAll_split_train.txt'
if __name__ == "__main__":
    if args.type == 'test':
        dataset = Test
        flnm = 'YLIMED_short_mp4_EvAll_split_test_dev.txt'
    elif args.type == 'trian':
        dataset = Train
        flnm = 'YLIMED_short_mp4_EvAll_split_train_dev.txt'

    data = pd.read_csv(dataset, sep=' ')
    X, y = data.ix[:, 0], data.ix[:, 2]

    # _, X_test, _, y_test = train_test_split(
    #     X, y, test_size=0.3, random_state=42, shuffle=True)
    # X_test = X_test.tolist()
    # y_test = y_test.tolist()

    s = StratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=0)
    s.get_n_splits(X, y)
    # X_test, y_test = X[0], y[]

    for train_index, test_index in s.split(X, y):
        print("TRAIN:", train_index, "TEST:", test_index)
        _, X_test = X[train_index], X[test_index]
        _, y_test = y[train_index], y[test_index]
        X_test = X_test.tolist()
        y_test = y_test.tolist()
        # print(X_test)
        with open(flnm, 'w') as f:
            for i in range(len(X_test)):
                f.write(' '.join([
                    str(X_test[i]),
                    str(X_test[i])[:5],
                    str(y_test[i])
                ]) + '\n')


    # with open('YLIMED_short_mp4_EvAll_split_train_dev.txt', 'w') as f:
    #     for i in range(len(X_test)):
    #         f.write(str(X_test[i]) + ' ' + str(y_test[i]) + '\n')

    # pd.merge()
